/*
 * Copyright 2017 - 2025 the original author or authors.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see [https://www.gnu.org/licenses/]
 */

package infra.test.util;

import org.hamcrest.Matcher;
import org.hamcrest.MatcherAssert;
import org.jspecify.annotations.Nullable;
import org.w3c.dom.Document;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;

import java.io.ByteArrayInputStream;
import java.util.Collections;
import java.util.Map;

import javax.xml.namespace.QName;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.xpath.XPath;
import javax.xml.xpath.XPathConstants;
import javax.xml.xpath.XPathExpression;
import javax.xml.xpath.XPathExpressionException;
import javax.xml.xpath.XPathFactory;

import infra.util.CollectionUtils;
import infra.util.StringUtils;
import infra.util.xml.SimpleNamespaceContext;

/**
 * A helper class for applying assertions via XPath expressions.
 *
 * @author Rossen Stoyanchev
 * @since 4.0
 */
public class XpathExpectationsHelper {

  private final String expression;

  private final XPathExpression xpathExpression;

  private final boolean hasNamespaces;

  /**
   * XpathExpectationsHelper constructor.
   *
   * @param expression the XPath expression
   * @param namespaces the XML namespaces referenced in the XPath expression, or {@code null}
   * @param args arguments to parameterize the XPath expression with using the
   * formatting specifiers defined in {@link String#format(String, Object...)}
   * @throws XPathExpressionException if expression compilation failed
   */
  public XpathExpectationsHelper(String expression, @Nullable Map<String, String> namespaces, Object... args)
          throws XPathExpressionException {

    this.expression = String.format(expression, args);
    this.xpathExpression = compileXpathExpression(this.expression, namespaces);
    this.hasNamespaces = CollectionUtils.isNotEmpty(namespaces);
  }

  private static XPathExpression compileXpathExpression(String expression,
          @Nullable Map<String, String> namespaces) throws XPathExpressionException {

    SimpleNamespaceContext namespaceContext = new SimpleNamespaceContext();
    namespaceContext.setBindings(namespaces != null ? namespaces : Collections.emptyMap());
    XPath xpath = XPathFactory.newInstance().newXPath();
    xpath.setNamespaceContext(namespaceContext);
    return xpath.compile(expression);
  }

  /**
   * Return the compiled XPath expression.
   */
  protected XPathExpression getXpathExpression() {
    return this.xpathExpression;
  }

  /**
   * Parse the content, evaluate the XPath expression as a {@link Node},
   * and assert it with the given {@code Matcher<Node>}.
   */
  public void assertNode(byte[] content, @Nullable String encoding, Matcher<? super Node> matcher)
          throws Exception {

    Node node = evaluateXpath(content, encoding, Node.class);
    MatcherAssert.assertThat("XPath " + this.expression, node, matcher);
  }

  /**
   * Parse the content, evaluate the XPath expression as a {@link NodeList},
   * and assert it with the given {@code Matcher<NodeList>}.
   *
   * @since 4.0
   */
  public void assertNodeList(byte[] content, @Nullable String encoding, Matcher<? super NodeList> matcher)
          throws Exception {

    Document document = parseXmlByteArray(content, encoding);
    NodeList nodeList = evaluateXpath(document, XPathConstants.NODESET, NodeList.class);
    MatcherAssert.assertThat("XPath " + this.getXpathExpression(), nodeList, matcher);
  }

  /**
   * Apply the XPath expression and assert the resulting content exists.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void exists(byte[] content, @Nullable String encoding) throws Exception {
    Node node = evaluateXpath(content, encoding, Node.class);
    AssertionErrors.assertNotNull("XPath " + this.expression + " does not exist", node);
  }

  /**
   * Apply the XPath expression and assert the resulting content does not exist.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void doesNotExist(byte[] content, @Nullable String encoding) throws Exception {
    Node node = evaluateXpath(content, encoding, Node.class);
    AssertionErrors.assertNull("XPath " + this.expression + " exists", node);
  }

  /**
   * Apply the XPath expression and assert the resulting content with the
   * given Hamcrest matcher.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertNodeCount(byte[] content, @Nullable String encoding, Matcher<? super Integer> matcher)
          throws Exception {

    NodeList nodeList = evaluateXpath(content, encoding, NodeList.class);
    String reason = "nodeCount for XPath " + this.expression;
    MatcherAssert.assertThat(reason, nodeList != null ? nodeList.getLength() : 0, matcher);
  }

  /**
   * Apply the XPath expression and assert the resulting content as an integer.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertNodeCount(byte[] content, @Nullable String encoding, int expectedCount) throws Exception {
    NodeList nodeList = evaluateXpath(content, encoding, NodeList.class);
    AssertionErrors.assertEquals("nodeCount for XPath " + this.expression, expectedCount,
            (nodeList != null ? nodeList.getLength() : 0));
  }

  /**
   * Apply the XPath expression and assert the resulting content with the
   * given Hamcrest matcher.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertString(byte[] content, @Nullable String encoding, Matcher<? super String> matcher)
          throws Exception {

    String actual = evaluateXpath(content, encoding, String.class);
    MatcherAssert.assertThat("XPath " + this.expression, actual, matcher);
  }

  /**
   * Apply the XPath expression and assert the resulting content as a String.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertString(byte[] content, @Nullable String encoding, String expectedValue) throws Exception {
    String actual = evaluateXpath(content, encoding, String.class);
    AssertionErrors.assertEquals("XPath " + this.expression, expectedValue, actual);
  }

  /**
   * Apply the XPath expression and assert the resulting content with the
   * given Hamcrest matcher.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertNumber(byte[] content, @Nullable String encoding, Matcher<? super Double> matcher) throws Exception {
    Double actual = evaluateXpath(content, encoding, Double.class);
    MatcherAssert.assertThat("XPath " + this.expression, actual, matcher);
  }

  /**
   * Apply the XPath expression and assert the resulting content as a Double.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertNumber(byte[] content, @Nullable String encoding, Double expectedValue) throws Exception {
    Double actual = evaluateXpath(content, encoding, Double.class);
    AssertionErrors.assertEquals("XPath " + this.expression, expectedValue, actual);
  }

  /**
   * Apply the XPath expression and assert the resulting content as a Boolean.
   *
   * @throws Exception if content parsing or expression evaluation fails
   */
  public void assertBoolean(byte[] content, @Nullable String encoding, boolean expectedValue) throws Exception {
    String actual = evaluateXpath(content, encoding, String.class);
    AssertionErrors.assertEquals("XPath " + this.expression, expectedValue, Boolean.parseBoolean(actual));
  }

  /**
   * Evaluate the XPath and return the resulting value.
   *
   * @param content the content to evaluate against
   * @param encoding the encoding to use (optionally)
   * @param targetClass the target class, one of Number, String, Boolean,
   * org.w3c.Node, or NodeList
   * @throws Exception if content parsing or expression evaluation fails
   * @since 4.0
   */
  @Nullable
  public <T> T evaluateXpath(byte[] content, @Nullable String encoding, Class<T> targetClass) throws Exception {
    Document document = parseXmlByteArray(content, encoding);
    return evaluateXpath(document, toQName(targetClass), targetClass);
  }

  /**
   * Parse the given XML content to a {@link Document}.
   *
   * @param xml the content to parse
   * @param encoding optional content encoding, if provided as metadata (e.g. in HTTP headers)
   * @return the parsed document
   */
  protected Document parseXmlByteArray(byte[] xml, @Nullable String encoding) throws Exception {
    DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
    factory.setNamespaceAware(this.hasNamespaces);
    DocumentBuilder documentBuilder = factory.newDocumentBuilder();
    InputSource inputSource = new InputSource(new ByteArrayInputStream(xml));
    if (StringUtils.hasText(encoding)) {
      inputSource.setEncoding(encoding);
    }
    return documentBuilder.parse(inputSource);
  }

  /**
   * Apply the XPath expression to given document.
   *
   * @throws XPathExpressionException if expression evaluation failed
   */
  @SuppressWarnings("unchecked")
  @Nullable
  protected <T> T evaluateXpath(Document document, QName evaluationType, Class<T> expectedClass)
          throws XPathExpressionException {

    return (T) getXpathExpression().evaluate(document, evaluationType);
  }

  private <T> QName toQName(Class<T> expectedClass) {
    QName evaluationType;
    if (Number.class.isAssignableFrom(expectedClass)) {
      evaluationType = XPathConstants.NUMBER;
    }
    else if (CharSequence.class.isAssignableFrom(expectedClass)) {
      evaluationType = XPathConstants.STRING;
    }
    else if (Boolean.class.isAssignableFrom(expectedClass)) {
      evaluationType = XPathConstants.BOOLEAN;
    }
    else if (Node.class.isAssignableFrom(expectedClass)) {
      evaluationType = XPathConstants.NODE;
    }
    else if (NodeList.class.isAssignableFrom(expectedClass)) {
      evaluationType = XPathConstants.NODESET;
    }
    else {
      throw new IllegalArgumentException("Unexpected target class " + expectedClass + ". " +
              "Supported: numbers, strings, boolean, and org.w3c.Node and NodeList");
    }
    return evaluationType;
  }

}
